import torch
import torch.nn as nn
from torchvision import models


class ResLstmNet(nn.Module):
    def __init__(self, num_classes=3, hidden_size=512, num_layers=2, dropout=0.5):
        super(ResLstmNet, self).__init__()
        resnet = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V1)
        self.features = nn.Sequential(*list(resnet.children())[:-1])
        for param in self.features.parameters():
            param.requires_grad = False

        self.fc = nn.Linear(2048, 512)

        self.lstm = nn.LSTM(
            input_size=512,
            hidden_size=hidden_size,
            num_layers=num_layers,
            batch_first=True,
            dropout=dropout
        )

        self.classifier = nn.Sequential(
            nn.Linear(hidden_size, 128),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(128, num_classes),
            # nn.Sigmoid()
        )

    def forward(self, x):
        batch_size, seq_len, c, h, w = x.shape
        x = x.view(batch_size*seq_len, c, h, w)
        features = self.features(x)
        features = features.view(batch_size, seq_len, -1)
        features = self.fc(features)
        features = features.view(batch_size, seq_len, -1)
        lstm_out, _ = self.lstm(features)
        lstm_output = lstm_out[:, -1, :]
        output = self.classifier(lstm_output)
        return output.squeeze(1)


if __name__ == '__main__':
    model = ResLstmNet()
